{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GC4aZuODw9Q9"
   },
   "source": [
    "# Notebook Setup \n",
    "The following cell will install Drake, checkout the manipulation repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  \n",
    "\n",
    "More details are available [here](http://manipulation.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lRBJuPiYOGgY"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import os, sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='14145ea656da5bdeb7a7ae4fbd7e28b8e464e544', drake_version='20201118', drake_build='nightly')\n",
    "\n",
    "from IPython import get_ipython\n",
    "from IPython.display import display, HTML, clear_output\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "import matplotlib.pyplot as plt, mpld3\n",
    "if running_as_notebook:\n",
    "  mpld3.enable_notebook()\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "from pydrake.all import (RigidTransform, RotationMatrix, BasicVector, LeafSystem, \n",
    "                        ConnectPlanarSceneGraphVisualizer, Simulator)\n",
    "import pydrake.multibody.jupyter_widgets\n",
    "from manipulation.utils import FindResource\n",
    "from copy import deepcopy\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.distributions.independent import Independent\n",
    "from scipy.signal import lfilter\n",
    "import random\n",
    "\n",
    "def AddPlanarBinAndManipuland(plant):\n",
    "  parser = pydrake.multibody.parsing.Parser(plant)\n",
    "  bin = parser.AddModelFromFile(FindResource(\"models/planar_bin.sdf\"))\n",
    "  plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"bin_base\", bin),\n",
    "                   RigidTransform(RotationMatrix.MakeZRotation(np.pi/2.0), [0,0,-0.015]))\n",
    "  planar_joint_frame = plant.AddFrame(\n",
    "    pydrake.multibody.tree.FixedOffsetFrame(\n",
    "      \"planar_joint_frame\",\n",
    "      plant.world_frame(),\n",
    "      RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
    "  box = parser.AddModelFromFile(FindResource(\"models/061_foam_brick.sdf\"))\n",
    "  box_frame = plant.AddFrame(\n",
    "    pydrake.multibody.tree.FixedOffsetFrame(\n",
    "      \"box_frame\", plant.GetFrameByName(\n",
    "      \"base_link\", box),\n",
    "      RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
    "  box_joint = plant.AddJoint(\n",
    "    pydrake.multibody.tree.PlanarJoint(\n",
    "    \"box_joint\", planar_joint_frame, box_frame))\n",
    "  box_joint.set_default_translation([0, 0.033400])\n",
    "\n",
    "\n",
    "def AddPlanarGripper(plant):\n",
    "  parser = pydrake.multibody.parsing.Parser(plant)\n",
    "  parser.package_map().Add(\n",
    "    \"wsg_50_description\",\n",
    "    os.path.dirname(\n",
    "      pydrake.common.FindResourceOrThrow(\n",
    "        \"drake/manipulation/models/wsg_50_description/package.xml\")))\n",
    "  gripper = parser.AddModelFromFile(\n",
    "    FindResource(\"models/schunk_wsg_50_welded_fingers.sdf\"), \"gripper\")\n",
    "  gripper_body = plant.GetBodyByName(\"body\", gripper)\n",
    "\n",
    "  # Add a planar joint the old fashioned way (so that I can have three actuators):\n",
    "  gripper_false_body1 = plant.AddRigidBody(\n",
    "    \"false_body1\", gripper,\n",
    "    pydrake.multibody.tree.SpatialInertia(\n",
    "      0, [0,0,0], pydrake.multibody.tree.UnitInertia(0,0,0)))\n",
    "  gripper_false_body2 = plant.AddRigidBody(\n",
    "    \"false_body2\", gripper,\n",
    "    pydrake.multibody.tree.SpatialInertia(\n",
    "      0, [0,0,0], pydrake.multibody.tree.UnitInertia(0,0,0)))\n",
    "  gripper_x = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\n",
    "    \"gripper_x\", plant.world_frame(), plant.GetFrameByName(\"false_body1\"), [1, 0, 0], -.2, .2))\n",
    "  plant.AddJointActuator(\"gripper_x\", gripper_x)\n",
    "  gripper_z = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\n",
    "    \"gripper_z\", plant.GetFrameByName(\"false_body1\"), plant.GetFrameByName(\"false_body2\"), [0, 0, 1], 0.0, 0.3))\n",
    "  gripper_z.set_default_translation(0.25)\n",
    "  plant.AddJointActuator(\"gripper_z\", gripper_z)\n",
    "  gripper_frame = plant.AddFrame(pydrake.multibody.tree.FixedOffsetFrame(\n",
    "    \"gripper_frame\", plant.GetFrameByName(\"body\", gripper),\n",
    "    RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
    "  gripper_theta = plant.AddJoint(pydrake.multibody.tree.RevoluteJoint(\n",
    "    \"gripper_theta\", plant.GetFrameByName(\"false_body2\"), gripper_frame, [0, -1, 0],  -np.pi/2, np.pi/2))\n",
    "  plant.AddJointActuator(\"gripper_theta\", gripper_theta)\n",
    "\n",
    "  return gripper\n",
    "\n",
    "def MakePlanarGripperOnlyPlant():\n",
    "  plant = pydrake.multibody.plant.MultibodyPlant(time_step=0.005)\n",
    "  AddPlanarGripper(plant)\n",
    "  plant.Finalize()\n",
    "  return plant\n",
    "\n",
    "\n",
    "class StepInput(LeafSystem):\n",
    "  \"\"\"provides a step-wise control interface\"\"\"\n",
    "  def __init__(self, u0=None):\n",
    "    LeafSystem.__init__(self)\n",
    "    self.DeclareVectorOutputPort('desired_state',\n",
    "                                 BasicVector(6),\n",
    "                                 self.DoCalcDesiredState)\n",
    "    self.u_input = np.zeros(6,)\n",
    "    if u0 is not None:\n",
    "      self.update_input(u0)\n",
    "    \n",
    "  def DoCalcDesiredState(self, context, output):\n",
    "    output.SetFromVector(self.u_input)\n",
    "\n",
    "  def update_input(self, u, mode='abs'):\n",
    "    if mode == 'abs':\n",
    "      self.u_input[0:3] = u\n",
    "    else:\n",
    "      self.u_input[0:3] = self.u_input[0:3] + u\n",
    "\n",
    "class PlanarGripperGym(object):\n",
    "  \"\"\"\n",
    "  a gym-like RL environment in Drake. Internally controlled\n",
    "  by forces/torques (inverse dynamics), the positions are the\n",
    "  desired positions to the controllers, not hardsetting the positions\n",
    "  in the simulator.\n",
    "  \"\"\"\n",
    "  def __init__(self,\n",
    "               render=True,\n",
    "               render_mode='drake',\n",
    "               reward_function=None):\n",
    "    builder = pydrake.systems.framework.DiagramBuilder()\n",
    "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(builder, time_step=0.01)\n",
    "    AddPlanarBinAndManipuland(plant)\n",
    "    gripper = AddPlanarGripper(plant)\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    controller_plant =MakePlanarGripperOnlyPlant()\n",
    "    N = controller_plant.num_positions()\n",
    "    kp = [100]*N\n",
    "    ki = [1]*N\n",
    "    kd = [2*np.sqrt(kp[0])]*N\n",
    "    controller = builder.AddSystem(pydrake.systems.controllers.InverseDynamicsController(controller_plant, kp, ki, kd, False))\n",
    "    builder.Connect(plant.get_state_output_port(gripper), controller.get_input_port_estimated_state())\n",
    "    builder.Connect(controller.get_output_port(), plant.get_actuation_input_port())\n",
    "\n",
    "    u0 = [0, 0.125, 0]\n",
    "    step_interface = StepInput(u0=u0)\n",
    "    builder.AddSystem(step_interface)\n",
    "    builder.Connect(step_interface.get_output_port(0), controller.get_input_port_desired_state())\n",
    "\n",
    "    # drake visualizer\n",
    "    if render:\n",
    "        self.vis = ConnectPlanarSceneGraphVisualizer(\n",
    "                    builder,\n",
    "                    scene_graph,\n",
    "                    output_port=scene_graph.GetOutputPort(\"lcm_visualization\"),\n",
    "                    xlim=[-0.3, 0.3],\n",
    "                    ylim=[-0.1, 0.5],\n",
    "                    show=False)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    simulator = Simulator(diagram, context)\n",
    "    simulator.Initialize()\n",
    "\n",
    "    # important to keep internal variables, otherwise segfault\n",
    "    # due to potential garbage collection\n",
    "    self.plant = plant\n",
    "    self.scene_graph = scene_graph\n",
    "    self.context = context\n",
    "    self.controller = controller\n",
    "    self.controller_plant = controller_plant\n",
    "    self.step_interface = step_interface\n",
    "    self.simulator = simulator\n",
    "    self.builder = builder\n",
    "    self.gripper = gripper\n",
    "    self.diagram = diagram\n",
    "    self.reward_function = reward_function\n",
    "\n",
    "    # save initial state\n",
    "    self.u0 = u0\n",
    "    self.reset()\n",
    "\n",
    "\n",
    "  def truncate_action(self, du_raw, dp_upper_bound=0.02, dtheta_upper_bound=0.5):\n",
    "    \"\"\"Truncate the action to make behaviors more natural\n",
    "\n",
    "    Args:\n",
    "        du_raw (numpy.ndarray): the raw control input, before being truncated\n",
    "        dp_upper_bound (float, optional): the maximum translation allowed in one action\n",
    "        dtheta_upper_bound (float, optional): the maximum rotation allowed in one action\n",
    "\n",
    "    Returns:\n",
    "        du (numpy.ndarray): the truncated action\n",
    "    The original action is modified in-place. \n",
    "    \"\"\"\n",
    "    du = deepcopy(du_raw) # ! prevent in-place modification\n",
    "    # bound movements\n",
    "    state = self.get_state()\n",
    "    x, y = state[3:5]\n",
    "    if x < -0.2 or x > 0.2 or y > 0.35 or y < 0.1:\n",
    "      # push gripper back towards the center of the env\n",
    "      du[0:2] = np.array([0, 0.25]) - state[0:2]\n",
    "    \n",
    "    theta = state[5]\n",
    "    unbounded_theta_next = theta + du[2]\n",
    "    # if unbounded_theta_next > np.pi or unbounded_theta_next < -np.pi/2:\n",
    "        # du[2] = 0.0\n",
    "    # truncate action\n",
    "    pos_mag = np.linalg.norm(du[0:2])\n",
    "    if pos_mag > dp_upper_bound:\n",
    "      du[0:2] = du[0:2] / pos_mag * dp_upper_bound\n",
    "    if abs(du[2]) > dtheta_upper_bound:\n",
    "      du[2] = dtheta_upper_bound * np.sign(du[2])\n",
    "    return du\n",
    "\n",
    "  def step(self, du, dt=0.2, dp_upper_bound=0.01, dtheta_upper_bound=0.1):\n",
    "    # limit the magnitude of the position change and angle change\n",
    "    self.step_interface.update_input(du, 'rel')\n",
    "    past_state = self.get_state()\n",
    "    try:\n",
    "        self.simulator.AdvanceTo(self.context.get_time() + dt)\n",
    "    except:\n",
    "        return\n",
    "\n",
    "    current_state = self.get_state()\n",
    "    reward, done, info = self.reward_function(current_state)\n",
    "    return self.get_state(), reward, done, info\n",
    "\n",
    "  def simulate(self, dt):\n",
    "    self.simulator.AdvanceTo(self.context.get_time() + dt)\n",
    "\n",
    "  def reset(self):\n",
    "    self.simulator.Initialize()\n",
    "    initial_pv = np.array([ 0., 0.,  0,  0.,\n",
    "        0.125,  0.0, 0.0,  0.0,\n",
    "        0.0,  0.0,  0.0,  0.0])\n",
    "    plant_context = self.plant.GetMyMutableContextFromRoot(self.context)\n",
    "    self.plant.SetPositionsAndVelocities(plant_context, initial_pv)\n",
    "    self.step_interface.update_input(self.u0, 'abs')\n",
    "    self.simulator.AdvanceTo(self.context.get_time() + 2.0)\n",
    "    return self.get_state()\n",
    "\n",
    "  def get_state(self, mode='p'):\n",
    "    \"\"\"\n",
    "    mode: 'p', 'v, 'pv'\n",
    "    \"\"\"\n",
    "    plant_context = self.plant.GetMyContextFromRoot(self.context)\n",
    "    pos_vel = self.plant.GetPositionsAndVelocities(plant_context)\n",
    "    if mode == 'p':\n",
    "      return deepcopy(pos_vel[0:6])\n",
    "    elif mode == 'v':\n",
    "      return deepcopy(pos_vel[6:])\n",
    "    else:\n",
    "      return deepcopy(pos_vel)\n",
    "\n",
    "def pad_to_last(nums, total_length, axis=-1, val=0):\n",
    "    \"\"\"Pad val to last in nums in given axis.\n",
    "\n",
    "    length of the result in given axis should be total_length.\n",
    "\n",
    "    Raises:\n",
    "      IndexError: If the input axis value is out of range of the nums array\n",
    "\n",
    "    Args:\n",
    "        nums (numpy.ndarray): The array to pad.\n",
    "        total_length (int): The final width of the Array.\n",
    "        axis (int): Axis along which a sum is performed.\n",
    "        val (int): The value to set the padded value.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Padded array\n",
    "\n",
    "    \"\"\"\n",
    "    tensor = torch.Tensor(nums)\n",
    "    axis = (axis + len(tensor.shape)) if axis < 0 else axis\n",
    "\n",
    "    if len(tensor.shape) <= axis:\n",
    "        raise IndexError('axis {} is out of range {}'.format(\n",
    "            axis, tensor.shape))\n",
    "\n",
    "    padding_config = [0, 0] * len(tensor.shape)\n",
    "    padding_idx = abs(axis - len(tensor.shape)) * 2 - 1\n",
    "    padding_config[padding_idx] = max(total_length - tensor.shape[axis], val)\n",
    "    return F.pad(tensor, padding_config)\n",
    "\n",
    "\n",
    "def filter_valids(tensor, valids):\n",
    "    \"\"\"Filter out tensor using valids (last index of valid tensors).\n",
    "\n",
    "    valids contains last indices of each rows.\n",
    "\n",
    "    Args:\n",
    "        tensor (torch.Tensor): The tensor to filter\n",
    "        valids (list[int]): Array of length of the valid values\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Filtered Tensor\n",
    "\n",
    "    \"\"\"\n",
    "    return [tensor[i][:valid] for i, valid in enumerate(valids)]\n",
    "\n",
    "def discount_cumsum(x, discount):\n",
    "    \"\"\"Discounted cumulative sum.\n",
    "\n",
    "    See https://docs.scipy.org/doc/scipy/reference/tutorial/signal.html#difference-equation-filtering  # noqa: E501\n",
    "    Here, we have y[t] - discount*y[t+1] = x[t]\n",
    "    or rev(y)[t] - discount*rev(y)[t-1] = rev(x)[t]\n",
    "\n",
    "    Args:\n",
    "        x (np.ndarrary): Input.\n",
    "        discount (float): Discount factor.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarrary: Discounted cumulative sum.\n",
    "\n",
    "\n",
    "    \"\"\"\n",
    "    return lfilter([1], [1, float(-discount)], x[::-1],\n",
    "                                axis=0)[::-1]\n",
    "\n",
    "class BatchDataset:\n",
    "    def __init__(self, inputs, batch_size, extra_inputs=None):\n",
    "        self._inputs = [i for i in inputs]\n",
    "        if extra_inputs is None:\n",
    "            extra_inputs = []\n",
    "        self._extra_inputs = extra_inputs\n",
    "        self._batch_size = batch_size\n",
    "        if batch_size is not None:\n",
    "            self._ids = np.arange(self._inputs[0].shape[0])\n",
    "            self.update()\n",
    "\n",
    "    @property\n",
    "    def number_batches(self):\n",
    "        if self._batch_size is None:\n",
    "            return 1\n",
    "        return int(np.ceil(self._inputs[0].shape[0] * 1.0 / self._batch_size))\n",
    "\n",
    "    def iterate(self, update=True):\n",
    "        if self._batch_size is None:\n",
    "            yield list(self._inputs) + list(self._extra_inputs)\n",
    "        else:\n",
    "            for itr in range(self.number_batches):\n",
    "                batch_start = itr * self._batch_size\n",
    "                batch_end = (itr + 1) * self._batch_size\n",
    "                batch_ids = self._ids[batch_start:batch_end]\n",
    "                batch = [d[batch_ids] for d in self._inputs]\n",
    "                yield list(batch) + list(self._extra_inputs)\n",
    "            if update:\n",
    "                self.update()\n",
    "\n",
    "    def update(self):\n",
    "        np.random.shuffle(self._ids)\n",
    "        \n",
    "class MLPGaussian(nn.Module):\n",
    "    def __init__(self, \n",
    "                input_dim=7,\n",
    "                output_dim=6,\n",
    "                num_hidden=2,\n",
    "                hidden_dim=128,\n",
    "                nonlinear_act=nn.ReLU,\n",
    "                hidden_w_init=nn.init.xavier_normal_,\n",
    "                hidden_b_init=nn.init.zeros_,\n",
    "                output_w_inits=nn.init.xavier_normal_,\n",
    "                output_b_inits=nn.init.zeros_):\n",
    "        super(MLPGaussian, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.fc_in = nn.Linear(input_dim, hidden_dim)\n",
    "        self.act = nonlinear_act()\n",
    "        self.fc_list = nn.ModuleList()\n",
    "        for i in range(num_hidden - 1):\n",
    "            hidden_layer = nn.Linear(hidden_dim, hidden_dim)\n",
    "            hidden_w_init(hidden_layer.weight)\n",
    "            hidden_b_init(hidden_layer.bias)\n",
    "            self.fc_list.append(hidden_layer)\n",
    "        \n",
    "        self.fc_out_mean = nn.Linear(hidden_dim, output_dim)\n",
    "        self.fc_out_var = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "        output_w_inits(self.fc_out_mean.weight)\n",
    "        output_b_inits(self.fc_out_mean.bias)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # out = x.view(-1, self.input_dim)\n",
    "        out = self.fc_in(x)\n",
    "        out = self.act(out)\n",
    "        for _, layer in enumerate(self.fc_list, start=0):\n",
    "            out = layer(out)\n",
    "            out = self.act(out)\n",
    "        out_mean = self.fc_out_mean(out)\n",
    "        out_var = self.fc_out_var(out)\n",
    "        out_var = self.act(out_var)\n",
    "        out_var = out_var + 0.001 # add a small bias to make sure it is not equal to 0\n",
    "        # return (out_mean, out_var)\n",
    "        return out_mean\n",
    "\n",
    "    \n",
    "def init_path():\n",
    "    path_sample = {'observations': [],\n",
    "                    'next_observations': [],\n",
    "                    'actions': [],\n",
    "                    'rewards': [],\n",
    "                    'infos': []}\n",
    "    return path_sample\n",
    "\n",
    "\n",
    "class REINFORCE:\n",
    "    def __init__(self,\n",
    "                 env,\n",
    "                 policy,\n",
    "                 value_function,\n",
    "                 util_compute_policy_loss,\n",
    "                 util_compute_value_loss,\n",
    "                 util_compute_advantage,\n",
    "                 use_advantage=False,\n",
    "                 gae_lambda=1.0,\n",
    "                 max_episode_length=100,\n",
    "                 discount_ratio=0.99,\n",
    "                 learning_rate=0.01):\n",
    "\n",
    "        self.env = env\n",
    "        self.max_episode_length = max_episode_length\n",
    "        self.discount = discount_ratio\n",
    "        self._gae_lambda = gae_lambda\n",
    "        self.util_compute_policy_loss = util_compute_policy_loss\n",
    "        self.util_compute_value_loss = util_compute_value_loss\n",
    "        self.util_compute_advantage = util_compute_advantage\n",
    "        self.use_advantage = use_advantage\n",
    "        \n",
    "        self.policy = policy\n",
    "        self._value_function = value_function\n",
    "        self._policy_optimizer = torch.optim.Adam(self.policy.parameters(),\n",
    "                                                  lr=learning_rate)\n",
    "        self._vf_optimizer = torch.optim.Adam(self._value_function.parameters(),\n",
    "                                                lr=learning_rate)\n",
    "        \n",
    "    \n",
    "    def _process_samples(self, paths):\n",
    "        r\"\"\"Process sample data based on the collected paths.\n",
    "\n",
    "        Notes: P is the maximum episode length (self.max_episode_length)\n",
    "\n",
    "        Args:\n",
    "            paths (list[dict]): A list of collected paths\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The observations of the environment\n",
    "                with shape :math:`(N, P, O*)`.\n",
    "            torch.Tensor: The actions fed to the environment\n",
    "                with shape :math:`(N, P, A*)`.\n",
    "            torch.Tensor: The acquired rewards with shape :math:`(N, P)`.\n",
    "            list[int]: Numbers of valid steps in each paths.\n",
    "            torch.Tensor: Value function estimation at each step\n",
    "                with shape :math:`(N, P)`.\n",
    "\n",
    "        \"\"\"\n",
    "        valids = torch.Tensor([len(path['actions']) for path in paths]).int()\n",
    "        obs = torch.stack([\n",
    "            pad_to_last(path['observations'],\n",
    "                        total_length=self.max_episode_length,\n",
    "                        axis=0) for path in paths\n",
    "        ])\n",
    "        actions = torch.stack([\n",
    "            pad_to_last(path['actions'],\n",
    "                        total_length=self.max_episode_length,\n",
    "                        axis=0) for path in paths\n",
    "        ])\n",
    "        rewards = torch.stack([\n",
    "            pad_to_last(path['rewards'], total_length=self.max_episode_length)\n",
    "            for path in paths\n",
    "        ])\n",
    "        returns = torch.stack([\n",
    "            pad_to_last(discount_cumsum(path['rewards'], self.discount).copy(),\n",
    "                        total_length=self.max_episode_length) for path in paths\n",
    "        ])\n",
    "        with torch.no_grad():\n",
    "            baselines = self._value_function(obs)\n",
    "\n",
    "        return obs, actions, rewards, returns, valids, baselines\n",
    "    \n",
    "    def collect_paths(self,\n",
    "                    process_action=None,\n",
    "                    batch_size=100,\n",
    "                    max_episode_length=100):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            env: simulation environment (e.g.Gym)\n",
    "            batch_size (int, optional): the number of episodes to be included in one patch.\n",
    "                                        Defaults to 100.\n",
    "            max_episode_length (int, optional): the maximum episode length\n",
    "        Returns:\n",
    "            [list]: a list of dicts, each dict stores data of one episode\n",
    "        \"\"\"\n",
    "        env = self.env\n",
    "        paths = []\n",
    "        invalid_episode_count = 0\n",
    "        while len(paths) < batch_size:\n",
    "            st= env.reset()\n",
    "            this_path = init_path()\n",
    "            good_episode = True\n",
    "\n",
    "            for t in range(max_episode_length):\n",
    "                # take a step\n",
    "                at, _ = self.policy.get_action(st)\n",
    "                if process_action is not None:\n",
    "                    at = process_action(at)\n",
    "                es = env.step(at)\n",
    "                if es is not None:\n",
    "                    st2, reward, done, info = es\n",
    "                    # env.render()\n",
    "                    this_path['observations'].append(st)\n",
    "                    this_path['next_observations'].append(st2)\n",
    "                    this_path['actions'].append(at)\n",
    "                    this_path['rewards'].append(reward)\n",
    "                    this_path['infos'].append(info)\n",
    "                else:\n",
    "                    ## discard this episode if the \n",
    "                    ## simulation/integration fails\n",
    "                    ## this is because the time step\n",
    "                    ## in our simulation\n",
    "                    ## plant's  is considerably \n",
    "                    ## large (0.01), in order to \n",
    "                    ## speed up simulation.\n",
    "                    good_episode = False\n",
    "                    invalid_episode_count += 1\n",
    "                    break\n",
    "\n",
    "                if done:\n",
    "                    break\n",
    "                st = st2\n",
    "            if good_episode:\n",
    "                for key in this_path.keys():\n",
    "                    this_path[key] = np.c_[this_path[key]]\n",
    "                this_path['rewards'] = this_path['rewards'].squeeze()\n",
    "                paths.append(this_path)\n",
    "        return paths\n",
    "    \n",
    "    def get_mini_batch_data(self, *inputs):\n",
    "        \"\"\"mini batch training data generator\"\"\"\n",
    "        batch_dataset = BatchDataset(inputs, None)\n",
    "        for _ in range(1):\n",
    "            for dataset in batch_dataset.iterate():\n",
    "                yield dataset\n",
    "\n",
    "    def train_from_episode_batch(self, itr, paths):\n",
    "        \"\"\"Train the algorithm from an episode batch\n",
    "\n",
    "        Args:\n",
    "            itr (int): Iteration number.\n",
    "            paths (list[dict]): A list of collected paths.\n",
    "\n",
    "        Returns:\n",
    "            numpy.float64: Calculated mean value of undiscounted returns.\n",
    "\n",
    "        \"\"\"\n",
    "        # the individual paths are not of the same length, \n",
    "        # the length of the paths are stored in valids\n",
    "        \n",
    "        obs, actions, rewards, returns, valids, baselines = \\\n",
    "            self._process_samples(paths)\n",
    "        obs_flat = torch.cat(filter_valids(obs, valids))\n",
    "        actions_flat = torch.cat(filter_valids(actions, valids))\n",
    "        rewards_flat = torch.cat(filter_valids(rewards, valids))\n",
    "        returns_flat = torch.cat(filter_valids(returns, valids))\n",
    "        advs_flat = self._compute_advantage(rewards, valids, baselines)\n",
    "\n",
    "        value_loss, policy_loss = 0., 0.\n",
    "        if self.use_advantage:\n",
    "            for dataset in self.get_mini_batch_data(\n",
    "                    obs_flat, actions_flat, advs_flat):\n",
    "                policy_loss = self._train_policy(*dataset)\n",
    "        else:\n",
    "            for dataset in self.get_mini_batch_data(\n",
    "                    obs_flat, actions_flat, returns_flat):\n",
    "                policy_loss = self._train_policy(*dataset)\n",
    "\n",
    "        for dataset in self.get_mini_batch_data(obs_flat, returns_flat):\n",
    "            value_loss = self._train_value_function(*dataset)\n",
    "\n",
    "        mean_episode_lengths = np.mean(valids.numpy())\n",
    "        undiscounted_returns = list(torch.sum(rewards, dim=1).numpy())\n",
    "\n",
    "        total_reward = float(float(rewards_flat.sum()))\n",
    "\n",
    "        print('{}, total reward:{}, policy_loss:{}, value_loss:{}, mean episode length:{}'.format(\n",
    "            itr, total_reward, policy_loss, value_loss, mean_episode_lengths\n",
    "        ))\n",
    "        \n",
    "        return total_reward, mean_episode_lengths, policy_loss, value_loss\n",
    "\n",
    "    def _train_policy(self, obs, actions, returns):\n",
    "        self._policy_optimizer.zero_grad()\n",
    "        loss = self.compute_policy_loss(obs, actions, returns)\n",
    "        loss.backward()\n",
    "        self._policy_optimizer.step()\n",
    "        return loss\n",
    "\n",
    "    def _train_value_function(self, obs, returns):\n",
    "        self._vf_optimizer.zero_grad()\n",
    "        loss = self.util_compute_value_loss(self._value_function, obs, returns)\n",
    "        loss.backward()\n",
    "        self._vf_optimizer.step()\n",
    "        return loss\n",
    "\n",
    "    def compute_policy_loss(self, obs, actions, returns):\n",
    "        \"\"\"call student's implementation\"\"\"\n",
    "        loss = self.util_compute_policy_loss(self.policy, obs, actions, returns)\n",
    "        return loss\n",
    "    \n",
    "\n",
    "    def _compute_advantage(self, rewards, valids, baselines):\n",
    "        \"\"\"compute the advantage function\n",
    "\n",
    "        Args:\n",
    "            rewards (torch.tensor): the episodic rewards, batch_size x max episode length\n",
    "            valids (torch.tensor): valid episodic lengths\n",
    "                    to allow concatenation of episodes of different lengths,\n",
    "                    by default, the max episode length is used instead of \n",
    "                    the actual episode lengths to construct reward and \n",
    "                    baseline tensors \n",
    "            baselines (torch.tensor): estimated baselines (state values) \n",
    "                                      batch_size x max episode length\n",
    "\n",
    "        Returns:\n",
    "            [torch.tensor]: flattened advantage function\n",
    "        \"\"\"\n",
    "        advantages = self.util_compute_advantage(self.discount, self._gae_lambda,\n",
    "                                        self.max_episode_length, baselines,\n",
    "                                        rewards)\n",
    "        advantage_flat = torch.cat(filter_valids(advantages, valids))\n",
    "\n",
    "        return advantage_flat\n",
    "\n",
    "def draw_training_stats(training_stats,\n",
    "                        title):\n",
    "    plt.rcParams.update({'font.size': 12})\n",
    "    fig, (ax1, ax2, ax3, ax4) = plt.subplots(nrows=4, ncols=1, sharex=True, figsize=(10,10))\n",
    "    axes = [ax1, ax2, ax3, ax4]\n",
    "    ax1.plot(training_stats['total_reward'], 'o-', label='total_rewards')\n",
    "    ax2.plot(training_stats['policy_loss'], 'o-', label='policy_loss')\n",
    "    ax3.plot(training_stats['value_loss'], 'o-', label='value_loss')\n",
    "    ax4.plot(training_stats['avg_episode_length'], 'o-', label='avg_episode_length')\n",
    "    ax1.set_title(title)\n",
    "    for ax in axes:\n",
    "        ax.legend()\n",
    "    axes[-1].set_xlabel('# of episodes')\n",
    "\n",
    "    ax3.set_yscale('log')\n",
    "    ax2.set_yscale('log')\n",
    "\n",
    "def visualize_policy(reward_function, policy=None):\n",
    "    env = PlanarGripperGym(render=True,\n",
    "                           reward_function=reward_function)\n",
    "    st = env.reset()\n",
    "    env.vis.start_recording()\n",
    "    for i in range(10):\n",
    "        st = env.get_state()\n",
    "        at = None\n",
    "        if policy is None:\n",
    "            at = np.random.rand(3,) * 0.1 - 0.05\n",
    "            at = env.truncate_action(at)\n",
    "        else:\n",
    "            at, _ = algo.policy.get_action(st)\n",
    "        at = env.truncate_action(at)\n",
    "        env.step(at)\n",
    "\n",
    "    env.vis.stop_recording()\n",
    "    ani = env.vis.get_recording_as_animation(repeat=True)\n",
    "    display(HTML(ani.to_jshtml()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BNUFdbcpOGgY"
   },
   "source": [
    "# REINFORCE\n",
    "In the lecture, you have been introduced to the idea of policy gradient, where the objective is to find a direct mapping from states to actions. In this exercise, you will implement REINFORCE, also commonly referred to as Vanilla Policy Gradient (VPG) method, to solve simple manipulation problems. To finish this exercise, you need to complete the 3 steps listed below:\n",
    "\n",
    "1. Implement the method to compute the loss of the policy function\n",
    "2. Implement the method to compute the loss of the value function\n",
    "3. Implement the method to compute the advantage function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q7aRTpchOGgY"
   },
   "source": [
    "Now let's check out the manipulation problem you are going to solve! The objective of the problem is to manipulate the foam block in the scene to maximize its $x$ coordinate. We intentionally make this problem extremely simple to help you train and debug faster. If you would like to solve a more difficult task, you are welcome to modify the reward function defined in the cell below. Note that since REINFORCE is an on-policy method, it is not known for being data efficient, so expect longer training time depending on how difficult your task is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 533
    },
    "id": "i9pqcLB3OGgY",
    "outputId": "afa1e2bf-ee34-4aaa-a880-a473fa281748"
   },
   "outputs": [],
   "source": [
    "def reward_function(observation):\n",
    "    \"\"\"\n",
    "    observation = [block_x, block_z, block_theta,\n",
    "                   gripper_x, gripper_z, gripper_theta]\n",
    "    \"\"\"\n",
    "    \n",
    "    # compute dense rewards\n",
    "    block_angle = observation[2]\n",
    "    gripper_angle = observation[5]\n",
    "\n",
    "    reward = observation[0]\n",
    "    done = False\n",
    "    if observation[0] > 0.1:\n",
    "        done = True\n",
    "        reward = 100. # reward goal completion\n",
    "    info = {} \n",
    "    return reward, done, info\n",
    "\n",
    "if running_as_notebook:\n",
    "    visualize_policy(reward_function, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i4t08PqPA3U9"
   },
   "source": [
    "In this exercise, we use the Markov Decision Process (MDP) formulation. We assume we have the full access to the poses of the gripper and foam block. In the following sections, we use states and observations interchangeably, both referring to the concatenation of block pose and gripper pose. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zzN6kklYOGga"
   },
   "source": [
    "## Review of REINFORCE Algorithm\n",
    "\n",
    "**REINFORCE** (Monte-Carlo policy gradient) relies on an estimated return by Monte-Carlo methods using episode samples to update the policy parameter $\\theta$. REINFORCE works because the expectation of the sample gradient is equal to the actual gradient:\n",
    "\n",
    "$$\n",
    "\\nabla_\\theta J(\\theta)  = \\mathbb{E}_{\\pi_\\theta} [\\nabla_\\theta \\ln \\pi_\\theta (x_t, u_t) Q^{\\pi_\\theta}(x_t, u_t) ]\n",
    "$$\n",
    "\n",
    "\n",
    "where $\\pi_\\theta$ is the policy parameterized by $\\theta$, and the action distribution computed from the policy is $\\pi_\\theta(u_t, x_t) = p_{\\pi_\\theta}(u_t \\vert x_t)$. Since $Q^{\\pi_\\theta}(x_t, u_t)$ is unknown, we instead use unbiased samples to approximate it \n",
    "\n",
    "$$ Q^\\pi(x_t, u_t) = \\mathbb{E}_\\pi[G_t \\vert x_t, u_t] $$\n",
    "\n",
    "where $G_t$ is the returns computed from samples\n",
    "\n",
    "\n",
    "Therefore we can measure $G_t$ from real sample trajectories and use that to update our policy gradient. It relies on a full trajectory and that\u2019s why it is a Monte-Carlo method.\n",
    "\n",
    "The REINFORCE algorithm is quite straightforward:\n",
    "\n",
    "Initialize the policy parameter $\\theta$ at random.\n",
    "\n",
    "Generate one trajectory on policy $\\pi_\\theta: x_1, u_1, r_2, x_2, u_2, \\dots, x_t$\n",
    "\n",
    "For t=1, 2, \u2026 , T:\n",
    "\n",
    "1. Estimate the the return $G_t= \\sum_{t=0}^{T} \\alpha^t R_t$\n",
    "\n",
    "2. Update policy parameters: $\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta \\ln \\pi_\\theta(x_t, u_t)G_t$, where $\\alpha$ is the learning rate\n",
    "\n",
    "\n",
    "A widely used variation of REINFORCE is to subtract a baseline value from the return $G_t$ to reduce the variance of gradient estimation while keeping the bias unchanged (Remember we always want to do this when possible). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N7QF_hV7OGga"
   },
   "source": [
    "To implement REINFORCE algorithm to solve continuous control tasks, you will inevitably need function approximators to estimate the policy distribution and optionally value function if an advantage function is used instead of the total discounted rewards for $G_t$. To help you quickly get to the core of the algorithm, we have provided you the implementation of the majority parts of the algorithm. You may find these codes in the `REINFORCE`, `PolicyEstimator`, `ValueEstimator` classes from setup cell. Note that for this exercise, you are not required to fully understand these codes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PufvcPKIOGga"
   },
   "outputs": [],
   "source": [
    "class PolicyEstimator(nn.Module):\n",
    "    def __init__(self,\n",
    "                 num_hidden,\n",
    "                 hidden_dim,\n",
    "                 obs_dim = None,\n",
    "                 action_dim = None,\n",
    "                 nonlinear_act=nn.Tanh,\n",
    "                 init_std = 1.0):\n",
    "        super(PolicyEstimator, self).__init__()\n",
    "        self._obs_dim = obs_dim\n",
    "        self._action_dim = action_dim\n",
    "\n",
    "        self.num_hidden = num_hidden\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.model = MLPGaussian(input_dim=self._obs_dim,\n",
    "                    output_dim=self._action_dim,\n",
    "                    num_hidden=num_hidden,\n",
    "                    hidden_dim=hidden_dim,\n",
    "                    nonlinear_act=nonlinear_act)\n",
    "\n",
    "        init_std_param = torch.Tensor([init_std]).log()\n",
    "        self._init_log_std = torch.nn.Parameter(init_std_param)\n",
    "\n",
    "    def predict(self, state):\n",
    "        state_tensor = state\n",
    "        if isinstance(state, np.ndarray):\n",
    "            state_tensor = torch.from_numpy(state).float().unsqueeze(0)\n",
    "        action_mean = self.model(state_tensor)\n",
    "        # post process the action mean\n",
    "        # action_mean = torch.sigmoid(action_mean)*0.005 - 0.0025\n",
    "        action_std = self._init_log_std.exp() * torch.ones_like(action_mean)\n",
    "        action_distribution = Normal(action_mean, action_std)\n",
    "        return Independent(action_distribution, 1)\n",
    "    \n",
    "    \n",
    "    def forward(self, observations):\n",
    "        dist = self.predict(observations)\n",
    "        return (dist, dict(mean=dist.mean,\n",
    "                log_std=(dist.variance**.5).log()))\n",
    "        # return dist\n",
    "    \n",
    "    def get_action(self, observation):\n",
    "        r\"\"\"Get a single action given an observation.\n",
    "\n",
    "        Args:\n",
    "            observation (np.ndarray): Observation from the environment.\n",
    "                Shape is :math:`env_spec.observation_space`.\n",
    "\n",
    "        Returns:\n",
    "            tuple:\n",
    "                * np.ndarray: Predicted action. Shape is\n",
    "                    :math:`env_spec.action_space`.\n",
    "                * dict:\n",
    "                    * np.ndarray[float]: Mean of the distribution\n",
    "                    * np.ndarray[float]: Standard deviation of logarithmic\n",
    "                        values of the distribution.\n",
    "        \"\"\"\n",
    "        if not isinstance(observation, np.ndarray) and not isinstance(\n",
    "                observation, torch.Tensor):\n",
    "            observation = self._env_spec.observation_space.flatten(observation)\n",
    "        with torch.no_grad():\n",
    "            if not isinstance(observation, torch.Tensor):\n",
    "                observation = torch.as_tensor(observation).float()\n",
    "            observation = observation.unsqueeze(0)\n",
    "            action, agent_infos = self.get_actions(observation)\n",
    "            return action[0], {k: v[0] for k, v in agent_infos.items()}\n",
    "\n",
    "    def get_actions(self, observations):\n",
    "        r\"\"\"Get actions given observations.\n",
    "\n",
    "        Args:\n",
    "            observations (np.ndarray): Observations from the environment.\n",
    "                Shape is :math:`batch_dim \\bullet env_spec.observation_space`.\n",
    "\n",
    "        Returns:\n",
    "            tuple:\n",
    "                * np.ndarray: Predicted actions.\n",
    "                    :math:`batch_dim \\bullet env_spec.action_space`.\n",
    "                * dict:\n",
    "                    * np.ndarray[float]: Mean of the distribution.\n",
    "                    * np.ndarray[float]: Standard deviation of logarithmic\n",
    "                        values of the distribution.\n",
    "        \"\"\"\n",
    "        if isinstance(observations, list):\n",
    "            if isinstance(observations[0], np.ndarray):\n",
    "                observations = np.stack(observations)\n",
    "            elif isinstance(observations[0], torch.Tensor):\n",
    "                observations = torch.stack(observations)\n",
    "        with torch.no_grad():\n",
    "            if not isinstance(observations, torch.Tensor):\n",
    "                observations = torch.as_tensor(observations).float().to()\n",
    "            dist, info = self.forward(observations)\n",
    "            return dist.sample().cpu().numpy(), {\n",
    "                k: v.detach().cpu().numpy()\n",
    "                for (k, v) in info.items()\n",
    "            }\n",
    "    def compute_log_likelihood(self, obs, actions):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            obs (torch.tensor): flattened observations\n",
    "            actions (torch.tensor): flattened actions\n",
    "        return \n",
    "            log_likelihoods (torch.tensor): log probabilities of \n",
    "                                          sampling the selected actions\n",
    "        \"\"\"\n",
    "        distribution = self.predict(obs)\n",
    "        return distribution.log_prob(actions)\n",
    "\n",
    "class ValueEstimator(nn.Module):\n",
    "    def __init__(self,\n",
    "                 num_hidden,\n",
    "                 hidden_dim,\n",
    "                 obs_dim = None,\n",
    "                 action_dim = None,\n",
    "                 nonlinear_act=nn.Tanh,\n",
    "                 init_std = 1.0):\n",
    "        super(ValueEstimator, self).__init__()\n",
    "        self._obs_dim = obs_dim\n",
    "        self._action_dim = action_dim\n",
    "        self.num_hidden = num_hidden\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.model = MLPGaussian(input_dim=self._obs_dim,\n",
    "                    output_dim=1,\n",
    "                    num_hidden=num_hidden,\n",
    "                    hidden_dim=hidden_dim,\n",
    "                    nonlinear_act=nonlinear_act)\n",
    "\n",
    "        init_std_param = torch.Tensor([init_std]).log()\n",
    "        self._init_log_std = torch.nn.Parameter(init_std_param)\n",
    "        self.module = self.model\n",
    "\n",
    "    def predict(self, state):\n",
    "        state_tensor = state\n",
    "        if isinstance(state, np.ndarray):\n",
    "            state_tensor = torch.from_numpy(state).float().unsqueeze(0)\n",
    "        value_mean = self.model(state_tensor)\n",
    "        value_std = self._init_log_std.exp() * torch.ones_like(value_mean)\n",
    "        value_distribution = Normal(value_mean, value_std)\n",
    "        return Independent(value_distribution, 1)\n",
    "    \n",
    "    def compute_log_likelihood(self, obs, returns):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            obs (torch.tensor): flattened observations\n",
    "            returns (torch.tensor): flattened a measure of gain\n",
    "                                    (e.g.advantage function or discounted rewards)\n",
    "        return \n",
    "            log_likelihoods (torch.tensor): log probabilities of \n",
    "                                          sampling the selected actions\n",
    "        \"\"\"\n",
    "        dist = self.predict(obs)\n",
    "        log_likelihoods = dist.log_prob(returns.reshape(-1, 1))\n",
    "        return log_likelihoods\n",
    "    \n",
    "    def forward(self, obs):\n",
    "        # torch.Size([246, 1000, 4])\n",
    "        dist = self.predict(obs)\n",
    "        return dist.mean.flatten(-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SeZ6Wm3EOGga"
   },
   "source": [
    "### (a) Compute Policy Loss\n",
    "\n",
    "From the previous section, as for training the policy model, we would like to perform gradient ascent \n",
    "\n",
    "$$\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta \\ln \\pi_\\theta(x_t, u_t)G_t$$\n",
    "\n",
    "Thanks to the auto-differentiation feature from PyTorch, we only need to compute the objective value of $ j_t^{\\pi} = \\ln \\pi_\\theta(x_t, u_t)G_t $. From there, simply following the procedure below will automatically compute gradients to optimize all relevant parameters.\n",
    "\n",
    "```\n",
    "optimizer.zero_grad()\n",
    "loss = compute_loss(...) # we are at this step\n",
    "loss.backward()\n",
    "optimizer.step() # by default, minimize loss\n",
    "```\n",
    "\n",
    "Evaluating objective values in batch can effectively reduce the variance in the computed gradients. \n",
    "\n",
    "$$ J_t^{\\pi} \\leftarrow \\frac{\\sum_{i=1}^N \\ln \\pi_\\theta(x_t, u_t)G_t}{N} $$\n",
    "\n",
    "where $\\pi_\\theta(x_t, u_t)$ can be obtained from the supplied method `policy.compute_action_log_likelihood`.\n",
    "\n",
    "By default, the PyTorch optimizer minimizes a scalar *loss*. Since in our case we would like to increase the log-likelihoods of generating better actions, the associated loss is \n",
    "$$ L_t^{\\pi} \\leftarrow - J_t^{\\pi}  $$\n",
    "\n",
    "Now in the cell below, implement the `util_compute_policy_loss(policy, obs, actions, returns)`. Note that `obs`, `actions`, `returns` are all pre-flattened. You may treat them as data in a batch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RGpvJHitOGga"
   },
   "outputs": [],
   "source": [
    "def util_compute_policy_loss(policy, obs, actions, returns):\n",
    "    \"\"\"compute policy loss given observations, actions, \n",
    "       and advantages (or total discounted rewards)\n",
    "       Note: One can either use the advantage function\n",
    "             or discounted returns\n",
    "\n",
    "    Args:\n",
    "        policy (PolicyEstimator): a policy instance\n",
    "        obs (torch.tensor): flattened observations Size: data_size, observation dim\n",
    "        actions (torch.tensor): flattened actions Size: data_size, action_dim\n",
    "        returns (torch.tensor): flattened total discounted returns. Size: data_size \n",
    "        advantages (Optional, torch.tensor): flattened advantages\n",
    "                                             advantages are not needed here\n",
    "\n",
    "    Returns:\n",
    "        torch.tensor: the mean policy loss, should have size of 1\n",
    "    \"\"\"\n",
    "    loss = torch.tensor([0.0])\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "heI_WdygPaEz"
   },
   "source": [
    "### (b) Compute Value Loss\n",
    "\n",
    "With only the policy loss computed above, you may already run the REINFORCE algorithm as you will see later. To reduce the variance of gradients further, another useful trick is to replace the discounted total rewards with an advantage function $A^{\\pi}$. In this case, the update rule is then\n",
    "\n",
    "$$\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta \\ln \\pi_\\theta(x_t, u_t)u_t^{\\pi}$$\n",
    "\n",
    "An advantage function intuitively compares the *advantage* of a particular action compared to an *average* action. The value of a particular station-action pair is $Q(x_t, u_t)$, whereas the value of an average action at state $x_t$ is $V(x_t)$. $V(x_t)$ is the value function of the Markov Decision Process (MDP). The estimated state values are also referred to as the baselines.\n",
    "\n",
    "Now our job is to learn a value function of the MDP problem. It should map from states to scalar values that estimate the average future rewards that can be received from the designated states. \n",
    "\n",
    "In the cell below, the returns, also the total discounted rewards, are provided to you. You should train the value function by \n",
    "$$ \\max_\\theta \\, \\ln p(V_t \\vert u_t, \\theta) $$\n",
    "where $\\theta$ here refers to the parameters of the value function, and $V_t$ is the training targets stored in `returns`.\n",
    "\n",
    "**Follow similar steps as in the policy loss, complete the `util_compute_value_loss(value_function, obs, returns)` below.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Mk93nxDqOGga"
   },
   "outputs": [],
   "source": [
    "def util_compute_value_loss(value_function, obs, returns):\n",
    "    \"\"\"compute value function loss given observations and returns\n",
    "       Note: One can either use the advantage function\n",
    "             or discounted returns\n",
    "\n",
    "    Args:\n",
    "        value_function (ValueEstimator): an instance of a value function\n",
    "        obs (torch.tensor): flattened observations\n",
    "        returns (torch.tensor): total discounted returns\n",
    "\n",
    "    Returns:\n",
    "        torch.tensor: the mean policy loss, should have size of 1\n",
    "    \"\"\"\n",
    "    loss = torch.tensor([0.0])\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nilLtBgVPkDg"
   },
   "source": [
    "### (c) Compute Advantages\n",
    "\n",
    "In this part, we will implement the computation of an advantage function.\n",
    "\n",
    "Mathematically, we can evaluate the *advantage* of a selected action $u_t$ at $x_t$ by\n",
    "\n",
    "$$\\delta_t = Q^\\pi(x_t, u_t) - V^\\pi(x_t) =  R(x_t) + \\gamma V'(x_{t+1}) - V'(x_t)$$\n",
    "\n",
    "where $V'(x_t), V'(x_{t+1})$ are evaluated by the value function learned above, $\\gamma$ is the discount factor.\n",
    "\n",
    "Note $\\delta_t$ is the advantage for a time step. For a trajectory of states and actions, we can compute the discounted sum of advantages similar to rewards.\n",
    "\n",
    "$$A = \\sum_{t=0}^T (\\lambda \\gamma)^t \\delta_t \\quad \\lambda \\in [0,1]$$\n",
    "\n",
    "**In the cell below, compute the advantage function by completing the `compute_advantages` method.**\n",
    "\n",
    "Note: the baselines are the estimated state values. The rewards are raw rewards, NOT the discounted sums. The shape of the rewards and baselines are both $N \\times T$. This is because an episode may terminate in fewer steps than the maximum episode length. The elements after the termination of episodes are set to 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CYKlkNrXOGgb"
   },
   "outputs": [],
   "source": [
    "def compute_advantages(discount, gae_lambda, max_episode_length, baselines,\n",
    "                       rewards):\n",
    "    \"\"\"Calculate advantages.\n",
    "\n",
    "    Advantages are a discounted cumulative sum.\n",
    "\n",
    "    Calculate advantages using a baseline according to Generalized Advantage\n",
    "    Estimation (GAE)\n",
    "    \n",
    "    Args:\n",
    "        discount (float): RL discount factor (i.e. gamma).\n",
    "        gae_lambda (float): Lambda, as used for Generalized Advantage\n",
    "            Estimation (GAE).\n",
    "        max_episode_length (int): Maximum length of a single episode.\n",
    "        baselines (torch.Tensor): A 2D vector of value function estimates with\n",
    "            shape (N, T), where N is the batch dimension (number of episodes)\n",
    "            and T is the maximum episode length experienced by the agent. If an\n",
    "            episode terminates in fewer than T time steps, the remaining\n",
    "            elements in that episode should be set to 0.\n",
    "        rewards (torch.Tensor): A 2D vector of per-step rewards with shape\n",
    "            (N, T), where N is the batch dimension (number of episodes) and T\n",
    "            is the maximum episode length experienced by the agent. If an\n",
    "            episode terminates in fewer than T time steps, the remaining\n",
    "            elements in that episode should be set to 0.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: A 2D vector of calculated advantage values with shape\n",
    "            (N, T), where N is the batch dimension (number of episodes) and T\n",
    "            is the maximum episode length experienced by the agent. If an\n",
    "            episode terminates in fewer than T time steps, the remaining values\n",
    "            in that episode should be set to 0.\n",
    "\n",
    "    \"\"\"\n",
    "    advantages = torch.zeros_like(rewards)\n",
    "    return advantages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dfZgOm1qOGgb"
   },
   "outputs": [],
   "source": [
    "# make everything deterministic\n",
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "env = PlanarGripperGym(render=False,\n",
    "                       reward_function=reward_function)\n",
    "action_dim = 3\n",
    "obs_dim = 6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kMi53IATR37n"
   },
   "source": [
    "### REINFORCE without Advantage Function\n",
    "\n",
    "Now we can compare the performance of the REINFORCE algorithm with and without using the advantage function.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 741
    },
    "id": "s43GiJqzOGgb",
    "outputId": "ca09b51a-9b7f-40e0-9a22-56d4f63e5b71"
   },
   "outputs": [],
   "source": [
    "policy = PolicyEstimator(action_dim=action_dim,\n",
    "                         obs_dim=obs_dim,\n",
    "                         num_hidden=2,\n",
    "                         hidden_dim=64,\n",
    "                         nonlinear_act=nn.ReLU)\n",
    "value_function = ValueEstimator(action_dim=action_dim,\n",
    "                                obs_dim=obs_dim,\n",
    "                                num_hidden=2,\n",
    "                                hidden_dim=64,\n",
    "                                nonlinear_act=nn.ReLU)\n",
    "use_advantage = False\n",
    "print('using advantage: {}'.format(use_advantage))\n",
    "algo = REINFORCE(env=env,\n",
    "                 policy = policy,\n",
    "                 value_function= value_function,\n",
    "                 util_compute_policy_loss = util_compute_policy_loss,\n",
    "                 util_compute_value_loss = util_compute_value_loss,\n",
    "                 util_compute_advantage = compute_advantages,\n",
    "                 use_advantage=use_advantage\n",
    "                )\n",
    "print('start training')\n",
    "training_stats = {'total_reward': [],\n",
    "                'policy_loss': [],\n",
    "                'value_loss': [],\n",
    "                'avg_episode_length': []}\n",
    "num_episodes = 50\n",
    "\n",
    "if running_as_notebook:\n",
    "    for i in range(num_episodes):\n",
    "        paths = algo.collect_paths(process_action=env.truncate_action,\n",
    "                                    batch_size=20,\n",
    "                                    max_episode_length=10)\n",
    "        total_reward, mean_episode_lengths, policy_loss, value_loss = \\\n",
    "                    algo.train_from_episode_batch(i, paths)\n",
    "\n",
    "        training_stats['total_reward'].append(total_reward)\n",
    "        training_stats['policy_loss'].append(policy_loss)\n",
    "        training_stats['value_loss'].append(value_loss)\n",
    "        training_stats['avg_episode_length'].append(mean_episode_lengths)\n",
    "    \n",
    "    clear_output()\n",
    "    draw_training_stats(training_stats, 'training_report')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 533
    },
    "id": "c5BruT_OOGgb",
    "outputId": "8e2b033d-ed83-42f4-e8b7-3b80ed6baa69"
   },
   "outputs": [],
   "source": [
    "if running_as_notebook:\n",
    "    visualize_policy(reward_function, policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ljWLRsFJR9nq"
   },
   "source": [
    "### REINFORCE with Advantage Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 741
    },
    "id": "g89Mq8r8OGgc",
    "outputId": "70861381-29aa-40a7-b194-58d51f53d961"
   },
   "outputs": [],
   "source": [
    "policy = PolicyEstimator(action_dim=action_dim,\n",
    "                         obs_dim=obs_dim,\n",
    "                         num_hidden=2,\n",
    "                         hidden_dim=64,\n",
    "                         nonlinear_act=nn.ReLU)\n",
    "value_function = ValueEstimator(action_dim=action_dim,\n",
    "                                obs_dim=obs_dim,\n",
    "                                num_hidden=2,\n",
    "                                hidden_dim=64,\n",
    "                                nonlinear_act=nn.ReLU)\n",
    "\n",
    "use_advantage = True\n",
    "print('using advantage: {}'.format(use_advantage))\n",
    "algo = REINFORCE(env=env,\n",
    "                 policy = policy,\n",
    "                 value_function= value_function,\n",
    "                 util_compute_policy_loss = util_compute_policy_loss,\n",
    "                 util_compute_value_loss = util_compute_value_loss,\n",
    "                 util_compute_advantage = compute_advantages,\n",
    "                 use_advantage=use_advantage\n",
    "                )\n",
    "print('start training')\n",
    "training_stats = {'total_reward': [],\n",
    "                'policy_loss': [],\n",
    "                'value_loss': [],\n",
    "                'avg_episode_length': []}\n",
    "num_episodes = 50\n",
    "\n",
    "if running_as_notebook:\n",
    "    for i in range(num_episodes):\n",
    "        paths = algo.collect_paths(process_action=env.truncate_action,\n",
    "                                    batch_size=20,\n",
    "                                    max_episode_length=10)\n",
    "        total_reward, mean_episode_lengths, policy_loss, value_loss = \\\n",
    "                    algo.train_from_episode_batch(i, paths)\n",
    "\n",
    "        training_stats['total_reward'].append(total_reward)\n",
    "        training_stats['policy_loss'].append(policy_loss)\n",
    "        training_stats['value_loss'].append(value_loss)\n",
    "        training_stats['avg_episode_length'].append(mean_episode_lengths)\n",
    "    \n",
    "    clear_output()\n",
    "    draw_training_stats(training_stats, 'training_report')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 533
    },
    "id": "Y1fe5TlpOGgc",
    "outputId": "bd662849-671c-479d-d8ce-a5aad3ed0f76"
   },
   "outputs": [],
   "source": [
    "if running_as_notebook:\n",
    "    visualize_policy(reward_function, policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LvTF2EBTB3Ur"
   },
   "source": [
    "Now if you have extra time, you may want to modify the reward function to learn more complicated skills. For example, you may learn the box flipping skill from the force control chapter using REINFORCE method!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3u1YKG0LOGgd"
   },
   "source": [
    "Weng, L. (2018, April 08). Policy Gradient Algorithms. Retrieved November 19, 2020, from https://lilianweng.github.io/lil-log/2018/04/08/policy-gradient-algorithms.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MwE8yNg58VQN"
   },
   "source": [
    "## How will this notebook be Graded?\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza. \n",
    "\n",
    "For submission of this assignment, you must do two things. \n",
    "- Download and submit the notebook `policy_gradient.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "- Write down your answers to 9.1.c to a separately pdf file and submit it to Gradescope's written submission section. \n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [2 pts] Correct implementation of `util_compute_policy_loss`\n",
    "- [2 pts] Correct implementation of `util_compute_value_loss`\n",
    "- [4 pts] Correct implementation of `compute_advantages`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vb_OFdG9Zt7K"
   },
   "outputs": [],
   "source": [
    "from manipulation.exercises.rl.test_vpg import TestVPG\n",
    "from manipulation.exercises.grader import Grader \n",
    "\n",
    "Grader.grade_output([TestVPG], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "celltoolbar": "Tags",
  "colab": {
   "collapsed_sections": [
    "GC4aZuODw9Q9"
   ],
   "name": "policy_gradient",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}